#pragma once
#include <memory>
#include <assert.h>
#include <iostream>
#include "math/checked_math.h"

/*
SmallVector是一个高效、适用于小规模数据的动态数组
它类似于标准库中的 `std::vector`，但在处理小规模数据时具有更好的性能，因为它避免了动态内存分配。
*/

namespace good_cptl
{

template<typename T, std::size_t N, typename A = std::allocator<T>>
class SmallVector
{
public:
  using value_type = T;
  using allocator_type = A;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;
  using reference = value_type&;
  using const_reference = const value_type&;
  using pointer = typename std::allocator_traits<A>::pointer;
  using const_pointer = typename std::allocator_traits<A>::const_pointer;
  using iterator = pointer;
  using const_iterator = const_pointer;

  SmallVector() = default;
  explicit SmallVector(size_t count, const T& value = T(), const A& alloc = A())
  : size_(0), capacity_(N), allocator_(alloc)
{
  std::cout << "count: " << count << std::endl;
  std::cout << "N: " << N << std::endl;
  if (count > N) { // 如果元素数量大于阈值，退化到堆内存
    heapStorage_ = allocator_.allocate(count);
    capacity_ = count;
  }
  for (size_type i = 0; i < count; ++i) {
    push_back(value);
  }
}

  SmallVector(const SmallVector& other);
  SmallVector(SmallVector&& other) noexcept;
  SmallVector& operator=(const SmallVector& other);
  SmallVector& operator=(SmallVector&& other) noexcept;

  ~SmallVector() { std::destroy(begin(), end());}

  // 元素访问
  reference operator[](size_type index) {
    assert(index >= 0 && index < size());
    return *(begin() + index);
  }
  const_reference operator[](size_type index) const {
    assert(index >= 0 && index < size());
    return *(begin() + index);
  }
  reference at(size_type index) { return operator[](index); }
  const_reference at(size_type index) const { return operator[](index); }
  pointer data() noexcept {
    if (is_using_heap()) {
      std::cout << "use heap " << std::endl;
      return heap();
    } else {
      std::cout << "use buffer " << std::endl;
      return buffer();
    }
  }
  const_pointer data() const noexcept {
    if (is_using_heap()) {
      return heap();
    } else {
      return buffer();
    }
  }
  reference front() {
    assert(!empty());
    return *begin();
  }
  const_reference front() const {
    assert(!empty());
    return *begin();
  }

  reference back() {
    assert(!empty());
    return *std::prev(end());
  }
  const_reference back() const {
    assert(!empty());
    return *std::prev(end());
  }

  // 容量相关接口
  bool empty() const noexcept { return size() == 0; }
  size_type size() const noexcept { return size_; }
  size_type capacity() const noexcept { return capacity_; }
  void reserve(size_type new_capacity);

  // 修改操作
  void clear() noexcept;

  void push_back(const T& value);
  void push_back(T&& value);
  template <class... Args>
  reference emplace_back(Args&&... args);
  void pop_back();
  void resize(size_type count);

  // 迭代器
  iterator begin() noexcept { return data(); }
  const_iterator begin() const noexcept { return data(); }
  iterator end() noexcept { return begin() + size_; }
  const_iterator end() const noexcept { return begin() + size_; }
  const_iterator cbegin() const noexcept { return data(); }
  const_iterator cend() const noexcept { return end(); }

  // 堆栈接口
  // 直接返回栈上的接口
  T* buffer() noexcept { return inlineStorage_; }
  const T* buffer() const noexcept { return const_cast<const value_type*>(this->inlineStorage_); }
  // 直接返回堆上的接口
  T* heap() noexcept { return heapStorage_; }
  const T* heap() const noexcept { return const_cast<const value_type*>(this->heapStorage_);}

private:
  void copy_inline_storage(const SmallVector& o);
  size_t compute_new_size() const;
  size_t max_size() const;
  void make_size(size_type new_size);
  template <typename EmplaceFunc>
  void make_size(size_type new_size, EmplaceFunc&& emplace_func, size_type pos);
  template <class EmplaceFunc>
  void move_to_uninitialized_emplace(T* begin, T* end, T* out,
                                     size_type pos, EmplaceFunc&& emplace_func);
  template <class EmplaceFunc>
  void make_size_internal(size_type new_size,
                          bool insert,
                          EmplaceFunc&& emplace_func,
                          size_type pos);

  void downsize(size_type sz);

  void move_to_uninitialized(T* first, T* last, T* out)
  {
    size_t idx = 0;
    for (; first != last; ++first, ++idx) {
      new(&out[idx]) T(std::move(*first));
    }
  }

  template <class Function>
  void populate_mem_forward(T* mem, std::size_t n, Function const& op)
  {
    std::size_t idx = 0;
    for (size_t i = 0; i < n; i++) {
      op(&mem[idx]);
      ++idx;
    }
  }



  // 计算新的内存

private:
   T inlineStorage_[N]; // 栈上的预分配空间
   T* heapStorage_;     // 堆上的分配空间

  bool is_using_heap() const noexcept {
    return capacity_ > N;
  }

  std::size_t size_;  // 元素数量
  std::size_t capacity_; // 容量
  A allocator_;  //  Allocator
};

template<typename T, std::size_t N, typename A>
SmallVector<T, N, A>::SmallVector(const SmallVector& o)
  : size_(0)
  , capacity_(N)
  , allocator_(std::allocator_traits<A>::select_on_container_copy_construction(o.allocator_))
{
  if (!is_using_heap() && o.size() < N) {
    copy_inline_storage(o);
  } else {
    auto n = o.size();
    make_size(n);
    std::uninitialized_copy(o.begin(), o.begin() + n, begin());
    size_ = n;
  }
}

template<typename T, std::size_t N, typename A>
SmallVector<T, N, A>::SmallVector(SmallVector&& other) noexcept
  : size_(other.size_)
  , capacity_(other.capacity_)
  , allocator_(std::move(other.allocator_))
{
  if (other.is_using_heap()) {
    heapStorage_ = other.heapStorage_;
    other.heapStorage_ = nullptr;
  } else {
    for (size_type i = 0; i < size_; i++) {
      new (inlineStorage_ + i) T(std::move(other.inlineStorage_[i]));
    }
  }
  other.size_ = 0;
  other.capacity_ = 0;
}

template<typename T, std::size_t N, typename A>
SmallVector<T, N, A>& good_cptl::SmallVector<T, N, A>::operator=(const SmallVector& o)
{
  if (this == &o) return *this;
  if (!is_using_heap() && o.size() < N) {
    copy_inline_storage(o);
  } else {
    auto n = o.size();
    make_size(n);
    std::uninitialized_copy(o.begin(), o.begin() + n, begin());
    size_ = n;
  }
  return *this;
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::reserve(size_type new_cap)
{
  make_size(new_cap);
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::resize(size_type sz)
{
  if (sz <= size()) {
    downsize(sz);
    return;
  }
  auto extra = sz - size();
  make_size(sz);
  populate_mem_forward(begin() + size(), extra, [&](void* p) {
    new (p) value_type();
  });
  size_ += extra;
}



template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::copy_inline_storage(const SmallVector& o)
{
  auto n = o.size();
  std::copy(o.buffer(), o.buffer() + n, buffer());
  size_ = o.size();
}

template<typename T, std::size_t N, typename A>
size_t SmallVector<T, N, A>::compute_new_size() const
{
  // 先乘 3，再除 2
  size_t c = capacity();
  if (!checked_mul(&c, c, size_t(3))) {
    throw "mul overflow";
  }
  c = c / 2 + 1;
  return static_cast<size_type>(std::min<size_t>(c, max_size()));
}

template<typename T, std::size_t N, typename A>
size_t SmallVector<T, N, A>::max_size() const
{
  if (is_using_heap()) {
    return N;
  } else {
    return 1000000 * sizeof(value_type);
  }
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::push_back(const T& value)
{
  emplace_back(value);
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::push_back(T&& value)
{
  emplace_back(std::move(value));
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::pop_back()
{
  downsize(size() - 1);
}


template<typename T, std::size_t N, typename A>
template <class... Args>
SmallVector<T, N, A>::reference SmallVector<T, N, A>::emplace_back(Args&&... args)
{
  if (size_ < N) {
    std::cout << "size: " << size_ << std::endl;
    new (buffer() + size_) value_type(std::forward<Args>(args)...);
    size_++;
    return *(buffer() + size_);
  }
  auto current_size = size_;
  auto current_capacity = capacity_;
  if (current_size == current_capacity) {
    make_size(current_size + 1,
              [&](void* p) { new (p) value_type(std::forward<Args>(args)...);},
              current_size);
  } else {
    new (heap() + current_size) value_type(std::forward<Args>(args)...);
  }
  size_++;
  return *(heap() + current_size);
}

template<typename T, std::size_t N, typename A>
template <class EmplaceFunc>
void SmallVector<T, N, A>::move_to_uninitialized_emplace(T* begin, T* end, T* out,
                                                         size_type pos, EmplaceFunc&& emplace_func)
{
  emplace_func(out + pos);
  if (begin) {
    move_to_uninitialized(begin, begin + pos, out);
  }
  if (begin + pos < end) {
    move_to_uninitialized(begin + pos, end, out + pos + 1);
  }
}

template<typename T, std::size_t N, typename A>
template <class EmplaceFunc>
void SmallVector<T, N, A>::make_size_internal(size_t new_size,
                                              bool insert,
                                              EmplaceFunc&& emplace_func,
                                              size_type pos)
{
  assert(new_size <= max_size());
  if (!is_using_heap()) {
    return;
  }
  new_size = std::fmax<size_t>(new_size, compute_new_size());
  size_t need_bytes = new_size;
  assert(!checked_mul(&need_bytes, need_bytes, sizeof(value_type)));
  size_t new_capacity = new_size;
  void* newh = std::malloc(new_capacity * sizeof(value_type));
  if (newh == nullptr) {
    throw std::bad_alloc();
  }
  value_type* newp = static_cast<value_type*>(newh);
  if (insert) {
    move_to_uninitialized_emplace(begin(), end(), newp, pos, std::forward<EmplaceFunc>(emplace_func));
  } else {
    if (data()) {
      move_to_uninitialized(begin(), end(), newp);
    }
  }
  std::destroy(begin(), end());
  heapStorage_ = newp;
  capacity_ = new_capacity;
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::make_size(size_type new_size)
{
  if (new_size <= capacity_) {
    return;
  }
  make_size_internal(new_size, false, [](void*){}, 0);
}

template<typename T, std::size_t N, typename A>
template <typename EmplaceFunc>
void SmallVector<T, N, A>::make_size(size_type new_size, EmplaceFunc&& emplace_func, size_type pos)
{
  if (new_size <= capacity_) {
    return;
  }
  make_size_internal(new_size, true, std::forward<EmplaceFunc>(emplace_func), pos);
}


template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::clear() noexcept
{
  downsize(0);
}

template<typename T, std::size_t N, typename A>
void SmallVector<T, N, A>::downsize(size_type sz)
{
  assert(sz <= size_);
  assert(sz <= capacity_);
  std::destroy(begin() + sz, end());
}

} // namespace good_cptl